import os
import re
import matplotlib.pyplot as plt

# Define regex patterns for parsing the data
energy_pattern = re.compile(r'FINAL_ETOT_IS ([\-0-9\.]+) eV')
magnetism_pattern = re.compile(r'Total Magnetism on atom:\s+Fe\s+([0-9\.]+)')
bond_length_pattern = re.compile(r'v(\d+\.\d+)')

# Initialize data storage
energies = {}
magnetisms = {}
bond_lengths = {}

# Read energy and bond lengths values
for dirname in os.listdir('.'):
    if os.path.isdir(dirname):
        # Extract the bond length from the directory name
        bond_length_match = bond_length_pattern.match(dirname)
        if bond_length_match:
            bond_length = float(bond_length_match.group(1))
            bond_lengths[dirname] = bond_length

        energy_file = os.path.join(dirname, 'OUT.bak', 'running_scf.log')
        if os.path.isfile(energy_file):
            with open(energy_file, 'r') as file:
                for line in file:
                    match = energy_pattern.search(line)
                    if match:
                        energies[dirname] = float(match.group(1))
                        break

# Read magnetism values from mulliken.txt files
# (This step can be removed if you don't need magnetism values)

# Read data from deltaspin if present
delta_spin_dir = 'deltaspin'
delta_spin_energy_file = os.path.join(delta_spin_dir, 'OUT.bak', 'running_scf.log')
if os.path.isfile(delta_spin_energy_file):
    with open(delta_spin_energy_file, 'r') as file:
        for line in file:
            match = energy_pattern.search(line)
            if match:
                energies[delta_spin_dir] = float(match.group(1))
                bond_lengths[delta_spin_dir] = 2.30  # Assuming deltaspin corresponds to bond length 2.30
                break

# Calculate relative energies with respect to the ground state
ground_state_energy = energies.get('v2.30')
relative_energies = {k: v - ground_state_energy for k, v in energies.items()}

# Prepare data for plotting
relative_energy_values = [relative_energies[k] for k in sorted(relative_energies.keys(), key=lambda x: bond_lengths[x])]
bond_length_values = [bond_lengths[k] for k in sorted(bond_lengths.keys(), key=lambda x: bond_lengths[x])]
labels = sorted(bond_lengths.keys())

# Plotting the data
plt.figure(figsize=(10, 5))
plt.plot(bond_length_values, relative_energy_values, 'ko-', label='without DeltaSpin')

# Highlight the deltaspin data point with a red circle
if delta_spin_dir in energies:
    plt.plot(bond_lengths[delta_spin_dir], relative_energies[delta_spin_dir], 'ro', label='Deltaspin')

plt.xlabel('Bond length (Angstrom)')
plt.ylabel('Relative Energy (eV)')
plt.title('Variation of Total Energy with Bond Length')
plt.legend()
plt.grid(True)

# Save the plot as a PNG file
plt.savefig('energy_vs_bond_length.png')
print("Plot saved as 'energy_vs_bond_length.png'.")
plt.show()
